﻿// CDlgMyAddImport.cpp: 实现文件
//

#include "pch.h"
#include "MY_PE_TOOL.h"
#include "CDlgMyAddImport.h"
#include "afxdialogex.h"


// CDlgMyAddImport 对话框

IMPLEMENT_DYNAMIC(CDlgMyAddImport, CDialogEx)

CDlgMyAddImport::CDlgMyAddImport(CWnd* pParent /*=nullptr*/)
	: CDialogEx(IDD_DIALOG5, pParent)
{

}

CDlgMyAddImport::~CDlgMyAddImport()
{
}

void CDlgMyAddImport::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_EDIT1, m_EdtDllName);
	DDX_Control(pDX, IDC_EDIT2, m_EdtAPIName);
	DDX_Control(pDX, IDC_LIST1, m_LstBoxName);
}


BEGIN_MESSAGE_MAP(CDlgMyAddImport, CDialogEx)
	ON_BN_CLICKED(IDC_BUTTON1, &CDlgMyAddImport::OnBnClickedButton1)
	ON_BN_CLICKED(IDC_BUTTON2, &CDlgMyAddImport::OnBnClickedButton2)
	ON_BN_CLICKED(IDC_BUTTON3, &CDlgMyAddImport::OnBnClickedButton3)
END_MESSAGE_MAP()


// CDlgMyAddImport 消息处理程序




void CDlgMyAddImport::SetFileData(char** pFileData)
{
	m_pFileData = pFileData;
}

void CDlgMyAddImport::SetFilePath(TCHAR* pFilePath)
{
	m_pFileFullPath = pFilePath;
}


//增加节
DWORD CDlgMyAddImport::AddSection(DWORD dwValue)
{
    DWORD dwWriteBytes = 0;
    DWORD dwSectionHeadSize = sizeof(IMAGE_SECTION_HEADER);

    PIMAGE_DOS_HEADER pDosHead = (PIMAGE_DOS_HEADER)(*m_pFileData);
    PIMAGE_NT_HEADERS pNtHead = (PIMAGE_NT_HEADERS)((DWORD)pDosHead + pDosHead->e_lfanew);
    PIMAGE_SECTION_HEADER pSectionHead = (PIMAGE_SECTION_HEADER)((DWORD)&pNtHead->OptionalHeader + pNtHead->FileHeader.SizeOfOptionalHeader);

    DWORD dwPEHeadRellySize = (DWORD)pSectionHead + dwSectionHeadSize * pNtHead->FileHeader.NumberOfSections;
    dwPEHeadRellySize -= (DWORD)pDosHead;

    if ((pNtHead->OptionalHeader.SizeOfHeaders - dwPEHeadRellySize) < dwSectionHeadSize)
    {
        MessageBox(_T("剩余的空间不足以增加一个节"), _T("提示"), MB_OK);
        return 0;
    }

    DWORD dwFileAli = CalcOffset(dwValue, pNtHead->OptionalHeader.FileAlignment);

    DWORD dwSectionIndex = pNtHead->FileHeader.NumberOfSections - 1;

    IMAGE_SECTION_HEADER SecHead = { 0 };

    SecHead.Misc.VirtualSize = dwFileAli;
    DWORD dwMemAli = CalcOffset(pSectionHead[dwSectionIndex].Misc.VirtualSize, pNtHead->OptionalHeader.SectionAlignment);
    SecHead.VirtualAddress = pSectionHead[dwSectionIndex].VirtualAddress + dwMemAli;
    SecHead.SizeOfRawData = dwFileAli;
    SecHead.PointerToRawData = pSectionHead[dwSectionIndex].PointerToRawData + pSectionHead[dwSectionIndex].SizeOfRawData;
    SecHead.Characteristics = IMAGE_SCN_MEM_WRITE | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_EXECUTE;

    //打开文件
    HANDLE hFile = CreateFile(m_pFileFullPath, GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

    if (hFile == INVALID_HANDLE_VALUE)
    {
        MessageBox(_T("打开文件失败"), _T("提示"), MB_OK);
        return 0;
    }


    //计算节表偏移
    DWORD dwost = (DWORD)&pNtHead->FileHeader.NumberOfSections - (DWORD)pDosHead;
    ::SetFilePointer(hFile, dwost, NULL, FILE_BEGIN);

    //写入节表项
    DWORD dwNewSec = pNtHead->FileHeader.NumberOfSections + 1;
    ::WriteFile(hFile, &dwNewSec, 4, &dwWriteBytes, NULL);

    //计算SizeofImage偏移
    dwost = (DWORD)&pNtHead->OptionalHeader.SizeOfImage - (DWORD)pDosHead;
    ::SetFilePointer(hFile, dwost, NULL, FILE_BEGIN);

    //写入节表项
    dwost = CalcOffset(dwValue, pNtHead->OptionalHeader.SectionAlignment) + pNtHead->OptionalHeader.SizeOfImage;
    ::WriteFile(hFile, &dwost, 4, &dwWriteBytes, NULL);

    //将节区数据写入节表项
    ::SetFilePointer(hFile, dwPEHeadRellySize, NULL, FILE_BEGIN);
    BOOL bRet = WriteFile(hFile, SecHead.Name, 8, &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.Misc.VirtualSize, sizeof(SecHead.Misc.VirtualSize), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.VirtualAddress, sizeof(SecHead.VirtualAddress), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.SizeOfRawData, sizeof(SecHead.SizeOfRawData), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.PointerToRawData, sizeof(SecHead.PointerToRawData), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.PointerToRelocations, sizeof(SecHead.PointerToRelocations), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.PointerToLinenumbers, sizeof(SecHead.PointerToLinenumbers), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.NumberOfRelocations, sizeof(SecHead.NumberOfRelocations), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.NumberOfLinenumbers, sizeof(SecHead.NumberOfLinenumbers), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.Characteristics, sizeof(SecHead.Characteristics), &dwWriteBytes, NULL);


    //写入节数据
    DWORD dwOffset = ::SetFilePointer(hFile, 0, NULL, FILE_END);
    char* pBuff = new char[dwFileAli];
    ZeroMemory(pBuff, dwFileAli);
    WriteFile(hFile, pBuff, dwFileAli, &dwWriteBytes, NULL);

    delete[] pBuff;
    //关闭文件
    CloseHandle(hFile);

    return dwOffset;
}

DWORD CDlgMyAddImport::CalcOffset(DWORD dwVal, DWORD dwAlign)
{
    if (dwVal % dwAlign == 0)
    {
        return dwVal;
    }
    else
    {
        return (dwVal / dwAlign + 1) * dwAlign;
    }

    return 0;
}

//计算增加的节需要多大空间
DWORD CDlgMyAddImport::AddSectionSize()
{
    PIMAGE_DOS_HEADER pDosHead = (PIMAGE_DOS_HEADER)(*m_pFileData);
    PIMAGE_NT_HEADERS pNtHead = (PIMAGE_NT_HEADERS)((DWORD)pDosHead + pDosHead->e_lfanew);

    DWORD dwSize = pNtHead->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size;

    CString strText;
    m_EdtDllName.GetWindowTextW(strText);

    dwSize = dwSize + strText.GetLength() + 1;

    int nCount = m_LstBoxName.GetCount();

    //遍历每一项，不能重复添加相同的函数名
    for (int i = 0; i < nCount; i++)
    {
        CString strListText;
        m_LstBoxName.GetText(i, strListText);

        dwSize = dwSize + strListText.GetLength() + 3;
    }

    dwSize = dwSize + 2 * 4 * (nCount + 1);

    return CalcOffset(dwSize, pNtHead->OptionalHeader.FileAlignment);
}


//构造导入表项，并写入文件
BOOL CDlgMyAddImport::CreateImportWriteFile(DWORD dwOffset, DWORD dwAddSecSize)
{
    /*
    导入表项 INT IAT DLL名称 API函数名称
    */
    DWORD dwImportOffsetValue = NEWFATORVA(dwOffset);

    //先申请空间
    char* pNewImportItem = new char[dwAddSecSize];
    ZeroMemory(pNewImportItem, dwAddSecSize);
    
    //先将旧的导入表项拷贝出来
    PIMAGE_DOS_HEADER pDosHead = (PIMAGE_DOS_HEADER)(*m_pFileData);
    PIMAGE_NT_HEADERS pNtHead = (PIMAGE_NT_HEADERS)((DWORD)pDosHead + pDosHead->e_lfanew);

    DWORD dwOldImportOffset = pNtHead->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
    DWORD dwSize = pNtHead->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size;

    DWORD dwFileOffset = RVATOFA(dwOldImportOffset);
    dwFileOffset += (DWORD)pDosHead;

    memcpy(pNewImportItem, (void *)dwFileOffset, dwSize);

    //先计算新导入表项的位置
    PIMAGE_IMPORT_DESCRIPTOR pNewImportAddr = (PIMAGE_IMPORT_DESCRIPTOR)(pNewImportItem + dwSize - sizeof(IMAGE_IMPORT_DESCRIPTOR));

    //空出INT 和 IAT的位置
    int nCount = m_LstBoxName.GetCount();
    DWORD dwIntAndIatSize = (nCount + 1) * 2 * 4;

    //先安排DLL名
    char* pDllName = pNewImportItem + dwSize + sizeof(IMAGE_IMPORT_DESCRIPTOR) + dwIntAndIatSize;
    CString strDllName;
    m_EdtDllName.GetWindowTextW(strDllName);
    CStringA strDllNameA(strDllName);
    memcpy(pDllName, strDllNameA, strDllNameA.GetLength() + 1);

    //计算dll名的RVA
    DWORD dwDllNameRva = NEWFATORVA(dwOffset + dwSize + sizeof(IMAGE_IMPORT_DESCRIPTOR) + dwIntAndIatSize);
    pNewImportAddr->Name = dwDllNameRva;

    //填写OriginalFirstThunk 和 FirstThunk
    DWORD dwOriginalFirstThunkRva = NEWFATORVA(dwOffset + dwSize + sizeof(IMAGE_IMPORT_DESCRIPTOR));
    pNewImportAddr->OriginalFirstThunk = dwOriginalFirstThunkRva;

    DWORD dwFirstThunkRva = NEWFATORVA(dwOffset + dwSize + sizeof(IMAGE_IMPORT_DESCRIPTOR) + (nCount + 1) * 4);
    pNewImportAddr->FirstThunk = dwFirstThunkRva;

    //获取API函数名 并且将INT和IAT填上值
    DWORD dwApiNameOffset = (DWORD)pDllName + strDllNameA.GetLength() + 1;
    for (int i = 0; i < nCount; i++)
    {
        CString strApiNameW;

        m_LstBoxName.GetText(i, strApiNameW);
        CStringA strApiNameA(strApiNameW);

        PIMAGE_IMPORT_BY_NAME pApiName = (PIMAGE_IMPORT_BY_NAME)dwApiNameOffset;
        pApiName->Hint = 0xcccc;
        memcpy(pApiName->Name, strApiNameA.GetBuffer(), strApiNameA.GetLength() + 1);

        DWORD *pAddressOfData = (DWORD *)(pNewImportItem + dwSize + sizeof(IMAGE_IMPORT_DESCRIPTOR) + i * 4);
        DWORD dwIntRva = NEWFATORVA(dwApiNameOffset - (DWORD)pNewImportItem + dwOffset);
        *pAddressOfData = dwIntRva;

        pAddressOfData = (DWORD *)((DWORD)pAddressOfData + (nCount + 1) * 4);
        *pAddressOfData = dwIntRva;

        dwApiNameOffset += strApiNameA.GetLength() + 1 + sizeof(pApiName->Hint);
    }

    //打开文件
    HANDLE hFile = CreateFile(m_pFileFullPath, GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

    if (hFile == INVALID_HANDLE_VALUE)
    {
        MessageBox(_T("打开文件失败"), _T("提示"), MB_OK);
        return FALSE;
    }

    //写入节数据
    DWORD dwWriteBytes = 0;

    //将导入表偏移修改
    DWORD dwImportOff = (DWORD)&pNtHead->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress - (DWORD)pDosHead;
    ::SetFilePointer(hFile, dwImportOff, NULL, FILE_BEGIN);
    WriteFile(hFile, &dwImportOffsetValue, sizeof(dwImportOffsetValue), &dwWriteBytes, NULL);

    //将数据写入文件
    ::SetFilePointer(hFile, dwOffset, NULL, FILE_BEGIN);
    WriteFile(hFile, pNewImportItem, dwAddSecSize, &dwWriteBytes, NULL);

    //关闭文件
    CloseHandle(hFile);
    delete[] pNewImportItem;
    return TRUE;
}


//根据RVA给出FA
DWORD CDlgMyAddImport::RVATOFA(DWORD dwRVA)
{
    PIMAGE_DOS_HEADER pDosHead = (PIMAGE_DOS_HEADER)(*m_pFileData);
    PIMAGE_NT_HEADERS pNtHead = (PIMAGE_NT_HEADERS)((DWORD)pDosHead + pDosHead->e_lfanew);
    PIMAGE_OPTIONAL_HEADER32 pOptionalHead = &pNtHead->OptionalHeader;
    PIMAGE_SECTION_HEADER pSectionHead = (PIMAGE_SECTION_HEADER)((DWORD)&pNtHead->OptionalHeader + pNtHead->FileHeader.SizeOfOptionalHeader);

    //在节表中
    for (int i = 0; i < pNtHead->FileHeader.NumberOfSections; i++)
    {
        //判断该RVA是否在这个节区中
        if (dwRVA >= pSectionHead[i].VirtualAddress && dwRVA < pSectionHead[i].VirtualAddress + pSectionHead[i].SizeOfRawData)
        {
            return (dwRVA - pSectionHead[i].VirtualAddress + pSectionHead[i].PointerToRawData);
        }
    }

    return 0;
}

//根据FA给出RVA
DWORD CDlgMyAddImport::FATORVA(DWORD dwFA)
{
    PIMAGE_DOS_HEADER pDosHead = (PIMAGE_DOS_HEADER)(*m_pFileData);
    PIMAGE_NT_HEADERS pNtHead = (PIMAGE_NT_HEADERS)((DWORD)pDosHead + pDosHead->e_lfanew);
    PIMAGE_OPTIONAL_HEADER32 pOptionalHead = &pNtHead->OptionalHeader;
    PIMAGE_SECTION_HEADER pSectionHead = (PIMAGE_SECTION_HEADER)((DWORD)&pNtHead->OptionalHeader + pNtHead->FileHeader.SizeOfOptionalHeader);

    //FA在节表中
    for (int i = 0; i < pNtHead->FileHeader.NumberOfSections; i++)
    {
        //判断该FA是否在这个节区中
        if (dwFA >= pSectionHead[i].PointerToRawData && dwFA < pSectionHead[i].PointerToRawData + pSectionHead[i].SizeOfRawData)
        {
            //如果在这个节区,计算RVA
            DWORD dwRVA = dwFA - pSectionHead[i].PointerToRawData + pSectionHead[i].VirtualAddress;

            return dwRVA;
        }
    }

    return 0;
}


//新数据FA转RVA
DWORD CDlgMyAddImport::NEWFATORVA(DWORD dwFA)
{
    //打开文件
    HANDLE hFile = CreateFile(m_pFileFullPath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

    if (hFile == INVALID_HANDLE_VALUE)
    {
        return FALSE;
    }


    //获取文件大小
    DWORD dwFileSize = GetFileSize(hFile, NULL);

    if (dwFileSize <= 0)
    {
        CloseHandle(hFile);
        return FALSE;
    }

    //申请空间
    DWORD dwBytesRead = 0;
    char *pNewFileData = new char[dwFileSize];
    ZeroMemory(pNewFileData, dwFileSize);

    //读文件
    ReadFile(hFile, pNewFileData, dwFileSize, &dwBytesRead, NULL);
    
    CloseHandle(hFile);

    PIMAGE_DOS_HEADER pDosHead = (PIMAGE_DOS_HEADER)(pNewFileData);
    PIMAGE_NT_HEADERS pNtHead = (PIMAGE_NT_HEADERS)((DWORD)pDosHead + pDosHead->e_lfanew);
    PIMAGE_OPTIONAL_HEADER32 pOptionalHead = &pNtHead->OptionalHeader;
    PIMAGE_SECTION_HEADER pSectionHead = (PIMAGE_SECTION_HEADER)((DWORD)&pNtHead->OptionalHeader + pNtHead->FileHeader.SizeOfOptionalHeader);

    //FA在节表中
    for (int i = 0; i < pNtHead->FileHeader.NumberOfSections; i++)
    {
        //判断该FA是否在这个节区中
        if (dwFA >= pSectionHead[i].PointerToRawData && dwFA < pSectionHead[i].PointerToRawData + pSectionHead[i].SizeOfRawData)
        {
            //如果在这个节区,计算RVA
            DWORD dwRVA = dwFA - pSectionHead[i].PointerToRawData + pSectionHead[i].VirtualAddress;

            return dwRVA;
        }
    }

    delete[] pNewFileData;
    return 0;
}

//点击确定按钮
void CDlgMyAddImport::OnBnClickedButton1()
{
	CString strText;
	m_EdtDllName.GetWindowTextW(strText);
	if (strText.IsEmpty())
	{
		MessageBox(_T("请输入Dll名称"), _T("提示"), MB_OK);
		return;
	}

    //获取ListBox有多少项
    int nCount = m_LstBoxName.GetCount();
	if (nCount == 0)
	{
        MessageBox(_T("增加失败,未检测到API名称"), _T("提示"), MB_OK);
        return;
	}

    //计算增加的节需要多大空间
    DWORD dwAddSecSize = AddSectionSize();

    //增加节
    DWORD dwOffset = AddSection(dwAddSecSize);
    if (dwOffset == 0)
    {
        return;
    }

    //根据输入的项构造导入表项并写入文件
    BOOL bRet = CreateImportWriteFile(dwOffset, dwAddSecSize);

    if (bRet)
    {
        MessageBox(_T("增加成功"), _T("提示"), MB_OK);
        PostMessageW(WM_CLOSE);
    }
}


//点击取消按钮
void CDlgMyAddImport::OnBnClickedButton2()
{
	PostMessageW(WM_CLOSE);
}


//点击 + 按钮
void CDlgMyAddImport::OnBnClickedButton3()
{
	CString strApiName;
	CString strListText;

	//获取输入的函数名
	m_EdtAPIName.GetWindowText(strApiName);

	//判断是否为空
	if (strApiName.IsEmpty())
	{
		return;
	}

	//获取ListBox有多少项
	int nCount = m_LstBoxName.GetCount();

	//遍历每一项，不能重复添加相同的函数名
	for (int i = 0; i < nCount; i++)
	{
		m_LstBoxName.GetText(i, strListText);
		if (strApiName == strListText)
		{
			MessageBox(_T("该API函数已经存在"), _T("提示"), MB_OK);
			return;
		}
	}

	//项ListBox中添加
	m_LstBoxName.AddString(strApiName);
}
